# -*- coding: utf-8 -*-
import os
import nltk
import logging
from pathlib import Path
from typing import Dict
from config import settings
from config.exceptions import NLTKResourceError

class NLTKLoader:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.resource_status: Dict[str, bool] = {}

        # 初始化NLTK数据路径
        self._setup_nltk_paths()

    def _setup_nltk_paths(self):
        """配置NLTK数据路径"""
        try:
            # 确保自定义路径存在
            settings.NLTK_DATA_PATH.mkdir(parents=True, exist_ok=True)

            # 添加自定义路径到NLTK搜索路径
            nltk.data.path.append(str(settings.NLTK_DATA_PATH))
            self.logger.info(f"NLTK数据路径已配置: {settings.NLTK_DATA_PATH}")
        except Exception as e:
            raise NLTKResourceError(f"路径配置失败: {str(e)}")

    def check_resources(self):
        """检查所有需要的NLTK资源"""
        self.logger.info("开始检查NLTK资源...")

        # 检查核心资源
        for resource, path in settings.NLTK_PACKAGES.items():
            self._check_single_resource(resource, path)

        # 检查自定义数据文件
        self._check_custom_datasets()

        self.logger.info("NLTK资源检查完成")

    def _check_single_resource(self, name: str, rel_path: str):
        """检查单个NLTK资源"""
        full_path = settings.NLTK_DATA_PATH / rel_path
        exists = nltk.data.find(f"corpora/{name}") is not None

        if exists:
            self.resource_status[name] = True
            self.logger.debug(f"资源 {name} 已存在: {full_path}")
        else:
            self.resource_status[name] = False
            self.logger.warning(f"资源 {name} 缺失: {full_path}")
            self._handle_missing_resource(name)

    def _handle_missing_resource(self, name: str):
        """处理缺失资源"""
        try:
            self.logger.info(f"尝试自动下载 {name}...")
            nltk.download(name, download_dir=str(settings.NLTK_DATA_PATH))
            self.resource_status[name] = True
            self.logger.info(f"成功下载 {name}")
        except Exception as e:
            self.logger.error(f"自动下载失败: {str(e)}")
            raise NLTKResourceError(
                f"请手动下载 {name} 并放置到 {settings.NLTK_DATA_PATH} 对应目录\n"
                f"官方下载地址: https://github.com/nltk/nltk_data/tree/gh-pages/packages"
            )

    def _check_custom_datasets(self):
        """检查自定义数据集（如不规则动词表）"""
        self.irregular_verbs = {}
        csv_path = settings.NLTK_DATA_PATH / "corpora" / "irregular_verbs.csv"

        try:
            if csv_path.exists():
                import pandas as pd
                df = pd.read_csv(csv_path)
                self.irregular_verbs = dict(zip(df['past'], df['base']))
                self.logger.info(f"加载自定义动词表: 共 {len(self.irregular_verbs)} 条记录")
            else:
                self.logger.warning(f"自定义动词表缺失: {csv_path}")
        except Exception as e:
            self.logger.error(f"自定义数据集加载失败: {str(e)}")
            raise NLTKResourceError("请检查CSV文件格式是否正确")

    def validate_environment(self):
        """验证NLTK环境是否就绪"""
        missing = [k for k, v in self.resource_status.items() if not v]
        if missing:
            raise NLTKResourceError(
                f"缺失关键NLTK资源: {', '.join(missing)}\n"
                f"请检查: {settings.NLTK_DATA_PATH}"
            )

if __name__ == "__main__":
    # 命令行测试
    import logging
    logging.basicConfig(level=logging.INFO)

    loader = NLTKLoader()
    loader.check_resources()
    loader.validate_environment()